import torch
from DEAD.Gradiant.AutoGrad import get_vector_field_component as gvc
from DEAD.Gradiant.AutoGrad import auto_grad

def cal_score_sdf(model,data):
    # 数据解包
    # latent_vector_index,lz和n_slots的维度是[batch_size]
    # bottom_theta_boundary 的维度是[batch_size, nrho_points, nz_points,3]
    # top_theta_boundary 的维度是[batch_size, nrho_points, nz_points,3]
    # burning_surface_coordinate,cavity_region_coordinate和solid_region_coordinate 的维度是[batch_size, n_vertices, 3]
    # cavity_region_sdf和solid_region_sdf 的维度是[batch_size, n_vertices,1]
    latent_vector_index, \
    burning_surface_coordinate, \
    cavity_region_coordinate, \
    cavity_region_sdf, \
    solid_region_coordinate, \
    solid_region_sdf = data

    # 设置梯度
    burning_surface_coordinate.requires_grad = False
    cavity_region_coordinate.requires_grad = False
    solid_region_coordinate.requires_grad = False
    with torch.no_grad():
        # 燃面边界条件项
        u_pred = model.forward_with_latent_vector_index(burning_surface_coordinate,latent_vector_index)
        burning_surface_term = torch.mean((u_pred)**2)
        
        # cavity区域项
        u_pred = model.forward_with_latent_vector_index(cavity_region_coordinate,latent_vector_index)
        cavity_term=torch.mean((u_pred-cavity_region_sdf)**2)

        # solid区域项
        u_pred = model.forward_with_latent_vector_index(solid_region_coordinate,latent_vector_index)
        solid_term=torch.mean((u_pred-solid_region_sdf)**2)

        score_value=(burning_surface_term+cavity_term+solid_term)/3.0
        score_dictionary={
        "burning_surface_term": burning_surface_term.item(),
        "cavity_term": cavity_term.item(),
        "solid_term": solid_term.item()
        }
    return score_value.item(), score_dictionary

def cal_score_eikonal(model,data):
    # 数据解包
    # latent_vector_index,lz和n_slots的维度是[batch_size]
    # bottom_theta_boundary 的维度是[batch_size, nrho_points, nz_points,3]
    # top_theta_boundary 的维度是[batch_size, nrho_points, nz_points,3]
    # burning_surface_coordinate,cavity_region_coordinate和solid_region_coordinate 的维度是[batch_size, n_vertices, 3]
    # cavity_region_sdf和solid_region_sdf 的维度是[batch_size, n_vertices,1]
    latent_vector_index, \
    burning_surface_coordinate, \
    cavity_region_coordinate, \
    cavity_region_sdf, \
    solid_region_coordinate, \
    solid_region_sdf= data

    # 设置梯度
    burning_surface_coordinate.requires_grad = True
    cavity_region_coordinate.requires_grad = True
    solid_region_coordinate.requires_grad = True

    ### u程函方程 1-ux^2-uy^2-uz^2=0
    # 燃面程函项
    u_pred = model.forward_with_latent_vector_index(burning_surface_coordinate,latent_vector_index)
    du_dxyz = auto_grad(u_pred, burning_surface_coordinate, 0, False)
    eikonal_surface_term = torch.mean((1.0-(gvc(du_dxyz, 0)**2+gvc(du_dxyz, 1)**2+gvc(du_dxyz, 2)**2))**2)
    
    # cavity区域项
    u_pred = model.forward_with_latent_vector_index(cavity_region_coordinate,latent_vector_index)
    du_dxyz = auto_grad(u_pred, cavity_region_coordinate, 0, False)
    eikonal_cavity_term = torch.mean((1.0-(gvc(du_dxyz, 0)**2+gvc(du_dxyz, 1)**2+gvc(du_dxyz, 2)**2))**2)
    
    # solid区域项
    u_pred = model.forward_with_latent_vector_index(solid_region_coordinate,latent_vector_index)
    du_dxyz = auto_grad(u_pred, solid_region_coordinate, 0, False)
    eikonal_solid_term = torch.mean((1.0-(gvc(du_dxyz, 0)**2+gvc(du_dxyz, 1)**2+gvc(du_dxyz, 2)**2))**2)
    
    score_value=(eikonal_surface_term+eikonal_cavity_term+eikonal_solid_term)/3.0
    score_dictionary={
    "eikonal_surface_term": eikonal_surface_term.item(),
    "eikonal_cavity_term": eikonal_cavity_term.item(),
    "eikonal_solid_term": eikonal_solid_term.item()
    }
    return score_value.item(), score_dictionary